
#include "cumsum_tiling.h"
#include "register/op_def_registry.h"

namespace optiling {
constexpr int32_t CORE_NUM = 1;
constexpr int32_t UB_SIZE = 196608;

constexpr int32_t X_INDEX = 0;
constexpr int32_t BYTE_BLOCK = 32;

constexpr int32_t BUF_COUNT = (1 + 1) * 2 + 1;

static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    CumsumTilingData tiling;
    uint32_t totalCoreNum = CORE_NUM;
    uint64_t ubSizePlatForm = UB_SIZE;

    const gert::Shape xShape = context->GetInputShape(X_INDEX)->GetStorageShape();
    const ge::DataType dataType = context->GetInputDesc(X_INDEX)->GetDataType();
    int32_t dTypeSize = ge::GetSizeByDataType(dataType);
    int32_t perBlockCount = BYTE_BLOCK / dTypeSize;
    uint64_t ubMaxProcCount = ubSizePlatForm / BUF_COUNT / BYTE_BLOCK * perBlockCount;
    tiling.set_ubMaxProcCount(ubMaxProcCount);

    context->SetTilingKey(101);
    if (dataType == ge::DT_FLOAT16) {
        context->SetTilingKey(201);
    } else if (dataType == ge::DT_INT32) {
        context->SetTilingKey(301);
    } else if (dataType == ge::DT_INT8) {
        context->SetTilingKey(401);
    }

    int64_t xDimCount = xShape.GetDimNum();
    tiling.set_xDimCount(xDimCount);

    uint64_t* dimArrayPtr = tiling.get_xDimArray();
    for (int64_t i = 0; i < xDimCount; i++) {
        dimArrayPtr[i] = xShape.GetDim(i);
    }

    auto attrs = context->GetAttrs();
    if (attrs == nullptr) {
        return ge::GRAPH_FAILED;
    }
    const bool* exclusivePtr = attrs->GetBool(0);
    tiling.set_exclusive(*exclusivePtr);
    const bool* reversePtr = attrs->GetBool(1);
    tiling.set_reverse(*reversePtr);

    context->SetBlockDim(totalCoreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    printf("==zf==Tiling Data, ubMaxProcCount:%lu, xDimCount:%lu, exclusive:%d, reverse:%d\n",
           tiling.get_ubMaxProcCount(), tiling.get_xDimCount(), tiling.get_exclusive(), tiling.get_reverse());

    // 以下代码是ascendebug必须的
    size_t* workspaces = context->GetWorkspaceSizes(1);
    workspaces[0] = 16 * 1024 * 1024;

    return ge::GRAPH_SUCCESS;
}
}  // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
    return GRAPH_SUCCESS;
}
}  // namespace ge

namespace ops {
class Cumsum : public OpDef {
public:
    explicit Cumsum(const char* name) : OpDef(name) {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("axis")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("exclusive").AttrType(OPTIONAL).Bool(false);
        this->Attr("reverse").AttrType(OPTIONAL).Bool(false);

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(Cumsum);
}  // namespace ops
